package com.wenbo.ocr.utils;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtSession;

import java.nio.FloatBuffer;
import java.util.Map;

/**
 * @Author wenbo
 * @Date 2023/10/18 22:06
 **/
public class ONNXRuntimeUtils  implements AutoCloseable{

    private OrtEnvironment env = OrtEnvironment.getEnvironment();

    public ONNXRuntimeUtils() {
    }

    public OrtSession.Result run(String modelPath, Map<String, OnnxTensor> inputs) {
        try {
            OrtSession session = this.env.createSession(modelPath);

            OrtSession.Result var4;
            try {
                var4 = session.run(inputs);
            } catch (Throwable var7) {
                if (session != null) {
                    try {
                        session.close();
                    } catch (Throwable var6) {
                        var7.addSuppressed(var6);
                    }
                }

                throw var7;
            }

            if (session != null) {
                session.close();
            }

            return var4;
        } catch (Exception var8) {
            LogUtils.printMessage("模型运行失败", var8, LogUtils.Level.ERROR);
            return null;
        }
    }

    public OnnxTensor createTensor(FloatBuffer data, long[] shape) {
        try {
            return OnnxTensor.createTensor(this.env, data, shape);
        } catch (Exception var4) {
            LogUtils.printMessage("创建张量失败", var4, LogUtils.Level.ERROR);
            return null;
        }
    }

    public OnnxTensor createTensor(float[] data, long[] shape) {
        return this.createTensor(FloatBuffer.wrap(data), shape);
    }

    public void close() throws Exception {
        this.env.close();
    }
}
